import torch
import matplotlib.pyplot as plt
from torch import nn

# ==================== 构建数据分布 ===================
x = torch.linspace(0, 2 * torch.pi, 100)
y = torch.sin(x) + 2
y += torch.normal(0, 0.2, y.shape)

plt.plot(torch.arange(0, len(y), 1).numpy(), y.detach().numpy(), "-", label="real")
# ==================== 设计序列模型数据 ===================
# 输入序列：连续的5条y的数据  输出序列：跟在前面5条数据之后的1条y的数据
# 总的来说，就是用5条数据去预测1条数据
# 注意：序列模型根本不关注x（时间）这一维度的信息
input_size = 5
output_size = 1
# 计算100条数据中有多少个连续的6条数据
steps = 100 - (input_size + output_size) + 1
# 定义所选的数据大小
data = torch.zeros((steps, input_size + output_size))  # (95,6)
# 通过顺序关系，将数据取出
for i in range(steps):
    data[i] = y[i:i + (input_size + output_size)]
inputs = data[:, :-1]  # (95,5)
outputs = data[:, -1:]  # (95,1)
# ==================== 设计序列模型 ===================
model = nn.Sequential(
    nn.Linear(5, 10),
    nn.Tanh(),
    nn.Linear(10, 1)
)
# ==================== 训练 ===================
# 损失函数
criterion = nn.MSELoss()
# 优化器
optimizer = torch.optim.SGD(model.parameters(), lr=0.01, momentum=0.9)
epochs = 1000
for epoch in range(epochs):
    optimizer.zero_grad()
    predicts = model(inputs)
    loss = criterion(predicts, outputs)
    loss.backward()
    optimizer.step()

    print(f"epoch:{epoch + 1}/{epochs} -- loss:{loss.item():.4f}")
# ==================== 序列预测 ===================
model.eval()
predicts = model(inputs)  # predicts (95,1)
plt.plot(torch.arange(0, len(predicts), 1).numpy(), predicts.detach().numpy(), "-", label="predict")
plt.legend()
plt.show()

# 问题：能不能输入为5，输出为5呢？
